import time

import numpy as np
import torch

from dm_control import viewer
from dm_env import TimeStep

from gym_aloha.env import AlohaEnv

from lerobot.common.envs.utils import preprocess_observation, check_env_attributes_and_types
from lerobot.common.policies.pretrained import PreTrainedPolicy
from lerobot.common.policies.utils import get_device_from_parameters


class ViewAlohaEnv(AlohaEnv):
    def __init__(self, task, obs_type="pixels", render_mode="rgb_array", observation_width=640, observation_height=480,
                 visualization_width=640, visualization_height=480):
        super().__init__(task, obs_type, render_mode, observation_width, observation_height, visualization_width,
                         visualization_height)

        self._policy: PreTrainedPolicy = None
        self._device: torch.device = None

    def generate_action(self, time_step: TimeStep):
        observation = self._format_raw_obs(time_step.observation)

        observation = preprocess_observation(observation)
        observation = {
            key: observation[key].to(self._device, non_blocking=self._device.type == "cuda") for key in observation
        }

        with torch.inference_mode():
            action = self._policy.select_action(observation)

        action = action.to("cpu").numpy()
        action = action.flatten()

        return action

    def set_policy(self, policy: PreTrainedPolicy):
        self._policy = policy
        self._device = get_device_from_parameters(self._policy)

    def view(self):
        viewer.launch(self._env, self.generate_action)
